#include <stdio.h>      /* for printf() and fprintf() */
#include <sys/socket.h> /* for socket(), bind(), and connect() */
#include <arpa/inet.h>  /* for sockaddr_in and inet_ntoa() */
#include <stdlib.h>     /* for atoi() and exit() */
#include <string.h>     /* for memset() */
#include <unistd.h>     /* for close() */

#define MAXPENDING 5    /* Maximum outstanding connection requests */

#define DNS_TABLE_LENGTH 32 /* Length of DNS table array */
#define URL_LENGTH 128  /* Length of URL array */
#define IP_LENGTH 32    /* Length of IP array */
#define RCVBUFSIZE 128   /* Size of receive buffer */

/* For logger */
#define DEBUG 0
#define INFO 1

void DieWithError(char *errorMessage);  /* Error handling function */
void HandleTCPClient(int clntSocket, void *dnsElement);   /* TCP client handling function */
void setupDnsTable(char *fileName, void *dnsElement);    /* Read DNS configuration from file */

/* Struct to hold dnsTable element */
typedef struct DnsElement {
    char domainName[URL_LENGTH];
    char ipAddr[IP_LENGTH];
}DnsElement;

int main(int argc, char *argv[]) {
    int servSock;                    /* Socket descriptor for server */
    int clntSock;                    /* Socket descriptor for client */
    struct sockaddr_in dnsServAddr; /* Local address */
    struct sockaddr_in dnsClntAddr; /* Client address */
    unsigned short dnsServPort;     /* Server port */
    unsigned int clntLen;            /* Length of client address data structure */

    /* Test for correct number of arguments */
    if (argc != 2) {
        fprintf(stderr, "Usage:  %s <Server Port>\n", argv[0]);
        exit(1);
    }

    /* Create DNS name resolution table */
    struct DnsElement dnsElement [DNS_TABLE_LENGTH];
    // Clear dnsElement to 0
    memset(&dnsElement, 0, sizeof(struct DnsElement)*DNS_TABLE_LENGTH);
	// Initialize DNS table with file content
    setupDnsTable("dns.conf", (void *)&dnsElement);

    dnsServPort = atoi(argv[1]);  /* First arg:  local port */

    /* Create socket for incoming connections */
    if ((servSock = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
        DieWithError("FATAL: socket() failed");
      
    /* Construct local address structure */
    memset(&dnsServAddr, 0, sizeof(dnsServAddr));   /* Zero out structure */
    dnsServAddr.sin_family = AF_INET;                /* Internet address family */
    dnsServAddr.sin_addr.s_addr = htonl(INADDR_ANY); /* Any incoming interface */
    dnsServAddr.sin_port = htons(dnsServPort);      /* Local port */

    /* Bind to the local address */
    if (bind(servSock, (struct sockaddr *) &dnsServAddr, sizeof(dnsServAddr)) < 0)
        DieWithError("FATAL: bind() failed");

    /* Mark the socket so it will listen for incoming connections */
    if (listen(servSock, MAXPENDING) < 0)
        DieWithError("FATAL: listen() failed");

    #if INFO == 1
    printf("INFO: Listening on port %d for incoming connection.\n", dnsServPort);
    #endif
    /* Run forever */
    for (;;) {
        /* Set the size of the in-out parameter */
        clntLen = sizeof(dnsClntAddr);

        /* Wait for a client to connect */
        if ((clntSock = accept(servSock, (struct sockaddr *) &dnsClntAddr, 
                               &clntLen)) < 0)
            DieWithError("FATAL: accept() failed");

        /* clntSock is connected to a client! */
        #if INFO == 1
        printf("INFO: Handling client %s\n", inet_ntoa(dnsClntAddr.sin_addr));
        #endif

        HandleTCPClient(clntSock, (void *)&dnsElement);
    }
    /* NOT REACHED */
    return 0;
}

void HandleTCPClient(int clntSocket, void *dnsElement) {
    char recvBuffer[RCVBUFSIZE];        /* Buffer for echo string */
    int recvMsgSize;                    /* Size of received message */
    struct DnsElement *dnsElementPtr = (struct DnsElement *) dnsElement; 

    // Clear echoBuffer to 0
    memset(recvBuffer, 0, sizeof(char)*RCVBUFSIZE);

    #if DEBUG == 1
    printf("DEBUG: Entering HandleTCPClient. echoBuffer=%s\n", recvBuffer);
    #endif
    /* Assemble the full message string */
    // Receive first message from client
    if ((recvMsgSize = recv(clntSocket, recvBuffer, RCVBUFSIZE-1, 0)) < 0)
        DieWithError("FATAL: recv() failed");
    recvBuffer[recvMsgSize] = '\0';     // NULL pad the string
    
    int recvBufferLength = strlen(recvBuffer);
    while (recvBuffer[strlen(recvBuffer)-1] != '\n') {
        recvBufferLength = strlen(recvBuffer);
        // See if there is more data to receive
        if ((recvMsgSize = recv(clntSocket, recvBuffer+recvMsgSize, (RCVBUFSIZE-1)-recvBufferLength, 0)) < 0)
            DieWithError("FATAL: recv() failed");
        recvBuffer[recvMsgSize+recvBufferLength] = '\0';
    }

    // Strip the new line character
    recvBuffer[strlen(recvBuffer)-1] = '\0';
    
    #if DEBUG == 1
    printf("DEBUG: Received full message from client. echoBuffer=%s\n", recvBuffer);
    #endif

    /* Compare URL stored in echoBuffer with DNS Table */
    int matchFound = 0;
    while ( strcmp(dnsElementPtr->domainName, "") != 0 ) {
        #if DEBUG == 1
        printf("DEBUG: Current dnsElementPtr->domainName is: %s\n", dnsElementPtr->domainName);
        #endif
        if (strcmp(dnsElementPtr->domainName, recvBuffer) == 0) {
            // If a match in URL is found, send the IP Address 
            //  associated with the URL back to client.
            if (send(clntSocket, dnsElementPtr->ipAddr, strlen(dnsElementPtr->ipAddr), 0) != strlen(dnsElementPtr->ipAddr)) {
                DieWithError("FATAL: send() failed");
            } else {
                matchFound = 1;
                break;
            }
        } else {
            dnsElementPtr++;    // Go to next dnsElement array index
        }
    }

    /* If no match is found, send the original string back. */
    if (matchFound == 0) {
        if (send(clntSocket, recvBuffer, strlen(recvBuffer), 0) != strlen(recvBuffer)) 
            DieWithError("FATAL: send() failed");
    }
    

    #if INFO == 1
    printf("INFO: Closing client socket.\n");
    #endif
    close(clntSocket);    /* Close client socket */
}

void DieWithError(char *errorMessage)
{
    perror(errorMessage);
    exit(1);
}

void setupDnsTable(char *fileName, void *dnsElement) {
    FILE *file;             // File pointer
    int lineCount = 0;      // Line of file
    char line[URL_LENGTH];  // String to hold line content
    struct DnsElement *dnsElementPtr = (struct DnsElement *) dnsElement; 

    #if INFO == 1
    printf("INFO: Opening %s to fill DNS table...\n", fileName);
    #endif
    // Open file for read
    if((file = fopen(fileName, "r")) == NULL) {
        perror("FATAL: Error Opening File.");
        exit(1);
    }
    /* Get each line from file */
    while( fgets(line, sizeof(line), file) != NULL ) {
        lineCount++;
        // Strip new line character
        line[strlen(line)-1] = '\0';
        if (lineCount%2 == 1) {
            // Setup domain name field
            strcpy(dnsElementPtr->domainName, line);
        } else {
            // Setup ip address field
            strcpy(dnsElementPtr->ipAddr, line);
            dnsElementPtr++;
        }
    }

    fclose(file);
}

